import streamlit as st
import os
import tempfile
from pathlib import Path
from langchain_community.chat_models import ChatOpenAI
from langchain.schema import SystemMessage, HumanMessage
from hls_utils import HLSOptimizer, HLSExecutor, HLSAnalyzer
from doc_loader import VitisHLSDocLoader
from project_generator import HLSProjectGenerator
import pandas as pd

# Set page config
st.set_page_config(
    page_title="Vitis HLS Intelligent Teaching Assistant",
    page_icon="🎓",
    layout="wide"
)

# Initialize session state
if 'code' not in st.session_state:
    st.session_state.code = ""
if 'testbench' not in st.session_state:
    st.session_state.testbench = ""
if 'optimization_history' not in st.session_state:
    st.session_state.optimization_history = []
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []
if 'llm' not in st.session_state:
    st.session_state.llm = None

def create_llm(
    model_name: str = "Llama-3.2-1B-Instruct-Hybrid",
    temperature: float = 0.0
) -> ChatOpenAI:
    """
    Create and initialize a ChatOpenAI client that connects to a local
    Lemonade Server via OpenAI-compatible HTTP API.

    Args:
        model_name (str): The identifier of the model hosted on Lemonade Server.
        temperature (float): Sampling temperature for responses (0.0 = deterministic).

    Returns:
        ChatOpenAI: An LLM client instance for generating chat completions.
    """
    try:
        # Base URL of the local Lemonade Server's OpenAI-compatible API
        base_url = "http://localhost:8000/api/v0"

        # Instantiate the ChatOpenAI client with local API settings
        llm = ChatOpenAI(
            model_name=model_name,           # ID of the model to use
            temperature=temperature,         # Controls randomness
            openai_api_base=base_url,        # Custom API endpoint
            openai_api_key="none",           # API key placeholder (not required)
            verbose=False                    # Disable verbose logging
        )
        return llm

    except Exception as e:
        # Display an error message in Streamlit if client creation fails
        st.error(f"Failed to create LLM client: {e}")
        return None

def qa_mode():
    """Q&A mode for HLS-related questions"""
    st.header("HLS Q&A Mode")
    
    if not st.session_state.llm:
        st.warning("Please initialize the LLM in the sidebar first!")
        return
    
    # Initialize document loader
    doc_loader = VitisHLSDocLoader()
    
    # Chat input
    user_question = st.text_input("Ask a question about Vitis HLS:")
    
    if user_question:
        # Get relevant context
        context = doc_loader.get_relevant_context(user_question)
        
        # Create system message with context
        system_prompt = f"""You are an expert in AMD Vitis HLS. Use the following documentation as reference:

{context}

Please provide a clear and concise answer to the user's question. If the documentation doesn't contain the answer, say so and provide your best knowledge about the topic."""

        # Get answer from LLM
        messages = [
            SystemMessage(content=system_prompt),
            HumanMessage(content=user_question)
        ]
        
        with st.spinner("Thinking..."):
            response = st.session_state.llm.invoke(messages)
            
            # Add to chat history
            st.session_state.chat_history.append({
                "question": user_question,
                "answer": response.content
            })
    
    # Display chat history
    for chat in st.session_state.chat_history:
        st.markdown("**Q:** " + chat["question"])
        st.markdown("**A:** " + chat["answer"])
        st.markdown("---")

def code_run_mode():
    """Run HLS using generated project folders"""
    st.title("Vitis HLS Project Runner")
    
    st.markdown("""
    Run HLS synthesis and simulation on generated project folders.
    This mode uses the project files generated in the Project Generator tab.
    """)
    
    # Check for available project folders
    project_folders = []
    
    # Look for extracted project folders in current directory
    current_dir = os.getcwd()
    if os.path.exists(current_dir):
        for item in os.listdir(current_dir):
            item_path = os.path.join(current_dir, item)
            if (os.path.isdir(item_path) and 
                item.endswith('_autogen') and 
                os.path.exists(os.path.join(item_path, 'run.py'))):
                project_folders.append(item_path)
    
    if not project_folders:
        st.warning("⚠️ No HLS project folders found!")
        st.info("""
        To use this mode:
        1. Go to the **Project Generator** tab
        2. Generate and download a project ZIP file
        3. Extract the ZIP file to your current directory
        4. Return to this tab to run the project
        """)
        return
    
    # Project selection
    st.subheader("📁 Select Project to Run")
    
    # Create a selectbox with just folder names for display
    folder_names = [os.path.basename(folder) for folder in project_folders]
    selected_folder_name = st.selectbox(
        "Choose a project folder:",
        folder_names,
        help="Select the HLS project you want to run"
    )
    
    # Get the full path of selected folder
    selected_folder = None
    for folder in project_folders:
        if os.path.basename(folder) == selected_folder_name:
            selected_folder = folder
            break
    
    if selected_folder:
        # Display project info
        st.subheader("📋 Project Information")
        
        col1, col2 = st.columns(2)
        with col1:
            st.write(f"**Project Path:** `{selected_folder}`")
            
        with col2:
            # Check for project files
            files_status = []
            required_files = ['run.py', 'hls_config.cfg']
            for file in required_files:
                file_path = os.path.join(selected_folder, file)
                exists = os.path.exists(file_path)
                status = "✅" if exists else "❌"
                files_status.append(f"{status} {file}")
            
            st.write("**Required Files:**")
            for status in files_status:
                st.write(status)
        
        # Check if reports already exist
        st.markdown("---")
        st.subheader("📊 Report Status")
        
        # Check for existing reports
        analyzer = HLSAnalyzer()
        existing_reports = analyzer.check_existing_reports(selected_folder)
        
        col1, col2 = st.columns(2)
        with col1:
            csim_status = "✅ Found" if existing_reports['csim_exists'] else "❌ Not found"
            st.write(f"**CSIM Report:** {csim_status}")
            if existing_reports['csim_exists']:
                st.caption(f"📍 {existing_reports['csim_path']}")
        
        with col2:
            syn_status = "✅ Found" if existing_reports['syn_exists'] else "❌ Not found"
            st.write(f"**Synthesis Report:** {syn_status}")
            if existing_reports['syn_exists']:
                st.caption(f"📍 {existing_reports['syn_path']}")
        
        # Show action buttons
        st.markdown("---")
        
        # If reports exist, show options
        if existing_reports['csim_exists'] or existing_reports['syn_exists']:
            col1, col2 = st.columns(2)
            with col1:
                show_existing = st.button("📋 Show Existing Reports", type="primary", use_container_width=True)
            with col2:
                force_rerun = st.button("🔄 Force Re-run HLS", use_container_width=True)
        else:
            # No reports exist, only show run button
            show_existing = False
            force_rerun = st.button("🚀 Run HLS Project", type="primary", use_container_width=True)
        
        # Handle button actions
        if show_existing and (existing_reports['csim_exists'] or existing_reports['syn_exists']):
            # Show existing reports without re-running
            with st.spinner("📊 Loading existing reports..."):
                # Create mock execution results for existing reports
                mock_results = {
                    'stdout': 'Reports loaded from existing files',
                    'stderr': ''
                }
                analysis_results = analyzer.analyze_project(mock_results['stdout'], mock_results['stderr'], selected_folder)
        
        elif force_rerun or (not existing_reports['csim_exists'] and not existing_reports['syn_exists']):
            # Run or re-run HLS project
            executor = HLSExecutor()
            
            try:
                with st.spinner("🔄 Running HLS synthesis and simulation..."):
                    # Execute HLS project
                    execution_results = executor.execute_project(selected_folder)
                    
                    # Analyze results
                    analysis_results = analyzer.analyze_project(execution_results['stdout'], execution_results['stderr'], selected_folder)
            except Exception as e:
                st.error(f"Error during execution: {str(e)}")
                st.error("Please check your project and try again")
                analysis_results = None
        else:
            # No action taken, don't show results
            analysis_results = None
        
        # Display results if available
        if 'analysis_results' in locals() and analysis_results is not None:
            st.markdown("---")
            st.subheader("📊 Analysis Results")
            
            # Display status
            if analysis_results['status'] == 'success':
                st.success("HLS execution completed successfully")
                
                # Display CSIM report
                st.subheader("🧪 CSIM Report")
                csim_report = analysis_results.get('csim_report', '')
                if csim_report.strip():
                    st.code(csim_report, language='text')
                else:
                    st.info("No CSIM report found or CSIM not run")
                
                # Display SYN report
                st.subheader("🔧 Synthesis Report")
                syn_report = analysis_results.get('syn_report', {})
                
                # Display Performance Estimates
                if 'performance' in syn_report and syn_report['performance'].strip():
                    st.write("**Performance Estimates**")
                    st.code(syn_report['performance'], language='text')
                else:
                    st.info("No performance estimates found")
                
                # Display Utilization Estimates
                if 'utilization' in syn_report and syn_report['utilization'].strip():
                    st.write("**Utilization Estimates**")
                    st.code(syn_report['utilization'], language='text')
                else:
                    st.info("No utilization estimates found")
                    
                # Show debug info if no reports found
                if not csim_report.strip() and not syn_report:
                    st.warning("⚠️ No HLS reports found. This might indicate:")
                    st.markdown("""
                    - The HLS run hasn't completed yet
                    - The HLS run failed
                    - Report files are in a different location
                    - Check the execution messages below for details
                    """)
            else:
                st.error("HLS execution failed")
            
            # Display messages
            if analysis_results['messages']:
                with st.expander("Execution Messages", expanded=False):
                    for message in analysis_results['messages']:
                        if message.strip():
                            st.text(message)

def project_generator_mode():
    """Generate complete HLS project for download"""
    st.title("HLS Project Generator")
    
    st.markdown("""
    Generate a complete HLS project package with all necessary files:
    - Source files (.cpp, .h)
    - Testbench files
    - Configuration files
    - Run scripts (Python, TCL, Shell)
    - Documentation (README, description.json)
    """)
    
    # Initialize session state for project generator
    if 'project_synthesis_code' not in st.session_state:
        st.session_state.project_synthesis_code = ""
    if 'project_testbench_code' not in st.session_state:
        st.session_state.project_testbench_code = ""
    if 'project_header_code' not in st.session_state:
        st.session_state.project_header_code = ""
    
    # Code input sections
    col1, col2 = st.columns([3, 1])
    
    with col1:
        st.subheader("🔧 HLS Synthesis Code")
        synthesis_code = st.text_area(
            "Enter your HLS C++ synthesis code:",
            height=250,
            value=st.session_state.project_synthesis_code,
            help="This will be the main synthesis file. The function name will be auto-detected."
        )
        
        st.subheader("📋 Header Code (Optional)")
        header_code = st.text_area(
            "Enter your header file code:",
            height=150,
            value=st.session_state.project_header_code,
            help="If not provided, a default header will be generated."
        )
        
        st.subheader("🧪 Testbench Code (Optional)")
        testbench_code = st.text_area(
            "Enter your testbench code:",
            height=200,
            value=st.session_state.project_testbench_code,
            help="If not provided, a default testbench will be generated."
        )
    
    with col2:
        st.subheader("⚙️ Project Settings")
        
        # Device part
        device_part = st.text_input(
            "Device Part:",
            value="xc7z020-clg484-1",
            placeholder="Enter FPGA device part (e.g., xc7z020-clg484-1)",
            help="Specify the target FPGA device part number"
        )
        
        # Clock period
        clock_period = st.number_input(
            "Clock Period (ns):",
            min_value=1,
            max_value=100,
            value=10,
            step=1,
            help="Specify the target clock period in nanoseconds (e.g., 10 ns = 100 MHz)"
        )
        
        # Additional options
        st.subheader("📦 Package Options")
        include_examples = st.checkbox("Include example files", value=True, help="Include example configuration files")
        include_docs = st.checkbox("Include documentation", value=True, help="Include README and description files")
        
        # Generate button
        st.markdown("---")
        generate_button = st.button("🚀 Generate Project Package", type="primary", use_container_width=True)
    
    # Generate project
    if generate_button:
        if not synthesis_code.strip():
            st.error("❌ Please enter synthesis code!")
            return
        
        try:
            with st.spinner("🏗️ Generating project package..."):
                # Create project generator
                generator = HLSProjectGenerator()
                
                # Generate project zip
                zip_path = generator.generate_project(
                    synthesis_code=synthesis_code,
                    testbench_code=testbench_code if testbench_code.strip() else None,
                    header_code=header_code if header_code.strip() else None,
                    device_part=device_part if device_part.strip() else "xc7z020-clg484-1",
                    clock_period=clock_period
                )
                
                # Extract top function name for display
                top_function = generator.extract_top_function(synthesis_code)
                project_name = f"{top_function}_autogen"
                
                # Success message
                st.success(f"✅ Project '{project_name}' generated successfully!")
                
                # File info
                file_size = os.path.getsize(zip_path)
                st.info(f"📁 Project folder created in current directory")
                st.info(f"📦 ZIP package size: {file_size / 1024:.1f} KB")
                
                # Download button
                with open(zip_path, 'rb') as f:
                    st.download_button(
                        label="📥 Download Project Package",
                        data=f.read(),
                        file_name=f"{project_name}.zip",
                        mime="application/zip",
                        use_container_width=True
                    )
                
                # Show file structure
                with st.expander("📂 View Project Structure", expanded=False):
                    st.code(f"""
{project_name}/
├── {top_function}.cpp              # Synthesis code
├── {top_function}.h                # Header file
├── {top_function}_test.cpp         # Testbench code
├── run.py                          # Python run script
├── hls_config.cfg                  # HLS configuration
├── description.json                # Project description
├── README                          # Documentation
├── run_hls.tcl                     # TCL run script
└── run_vitis_commandline.sh        # Shell run script
                    """, language="text")
                
                # Store in session state
                st.session_state.project_synthesis_code = synthesis_code
                st.session_state.project_testbench_code = testbench_code
                st.session_state.project_header_code = header_code
                
                # Clean up temporary file
                try:
                    os.remove(zip_path)
                except:
                    pass
                
        except Exception as e:
            st.error(f"❌ Error generating project: {str(e)}")
            st.error("Please check your code and try again.")
    
    # Help section
    with st.expander("ℹ️ Help & Examples", expanded=False):
        st.markdown("""
        ### How to use:
        1. **Enter your synthesis code** - This should contain your main HLS function
        2. **Optionally add header and testbench code** - If not provided, defaults will be generated
        3. **Specify device part** - Choose the target FPGA device (e.g., xc7z020-clg484-1)
        4. **Click Generate** - A complete project package will be created
        5. **Download the ZIP file** - Contains all files needed to run your HLS project
        
        ### What's included:
        - **Source files**: Your code exactly as entered
        - **Configuration files**: HLS settings with your specified device part
        - **Run scripts**: Python, TCL, and shell scripts for different execution methods
        - **Documentation**: README and project description files

        ### Running the Design
                    
        - **Running the Design using Vitis HLS **(edit run_hls.tcl to set $hls_exec and enable specific run steps)
        ```
        vitis-run --mode hls --tcl run_hls.tcl
        ```

        - **Running the design using Vitis v++ and vitis-run commands**
        ```
        ./run_vitis_commandline.sh
        ```

        - **Running the design using python script**
        ```
        vitis -s run.py
        ```
    
        """)

def main():
    st.title("Vitis HLS Intelligent Teaching Assistant")
    st.markdown("""
    This tool helps you optimize your HLS C++ code through:
    - Code analysis and optimization suggestions
    - Automatic pragma insertion
    - HLS execution and performance analysis
    - Detailed optimization reports
    """)

    # Sidebar for settings
    with st.sidebar:
        st.header("Settings")
        model_name = st.selectbox(
            "Select Model",
            ["Llama-3.2-1B-Instruct-Hybrid", 
            "Llama-3.2-3B-Instruct-Hybrid",
            "Phi-3-Mini-Instruct-Hybrid",
            "Phi-3.5-Mini-Instruct-Hybrid",
            "Qwen-1.5-7B-Chat-Hybrid",
            "DeepSeek-R1-Distill-Llama-8B-Hybrid",
            "DeepSeek-R1-Distill-Qwen-7B-Hybrid",
            "Mistral-7B-v0.3-Instruct-Hybrid",
            "Llama-3.1-8B-Instruct-Hybrid",
            "Qwen3-8B-GGUF",
            "DeepSeek-Qwen3-8B-GGUF"],
            index=0
        )
        temperature = st.slider("Temperature", 0.0, 1.0, 0.0, 0.1)
        
        if st.button("Initialize LLM"):
            st.session_state.llm = create_llm(model_name, temperature)
            if st.session_state.llm:
                st.success("LLM initialized successfully!")

    # Create tabs for different modes
    tab1, tab2, tab3, tab4 = st.tabs(["Q&A Mode", "Project Generator", "RUN Mode", "Optimization Mode"])
    
    with tab1:
        qa_mode()
    
    with tab2:
        project_generator_mode()
        
    with tab3:
        code_run_mode()
        
    with tab4:
        st.info("🚧 Optimization Mode - Coming Soon!")
        st.markdown("This mode will provide AI-powered HLS code optimization suggestions.")

if __name__ == "__main__":
    main() 